from torch import (
    Tensor,
    tensor,
    isnan,
    isinf
)
from torch.nn import (
    Module,
    Embedding,
    Dropout,
    Linear,
    LSTM
)
from torchcrf import CRF
from utils import conf, data

class MovieBiLSTM_CRF(Module):
    def __init__(self):
        super(MovieBiLSTM_CRF, self).__init__()
        # Get data and word2id
        _, self.word2id = data.build_vocab()
        # Get model configuration
        self.embedding_size = conf.model.embedding_size
        self.hidden_size = conf.model.hidden_size
        self.vocab_size = len(self.word2id)
        self.labels = data.labels
        self.tag_size = len(self.labels)
        self._dropout = conf.model.dropout

        # Initialize word embedding
        self.word_embedding = Embedding(
            self.vocab_size,
            self.embedding_size
        )
        # Initialize LSTM
        self.lstm = LSTM(
            self.embedding_size,
            self.hidden_size // 2,
            bidirectional=True,
            batch_first=True
        )
        # Initialize dropout
        self.dropout = Dropout(self._dropout)
        # Initialize linear
        self.linear = Linear(
            self.hidden_size,
            self.tag_size
        )
        # Initialize CRF
        self.crf = CRF(
            self.tag_size,
            batch_first=True
        )

    def get_lstm2linear(self, x: Tensor) -> Tensor:
        # Get word embedding
        embed = self.word_embedding(x)
        # Get LSTM output
        lstm_out, _ = self.lstm(embed)
        # Apply dropout
        out = self.dropout(lstm_out)
        # Get linear output
        out = self.linear(out)
        return out

    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
        # Get LSTM output
        out = self.get_lstm2linear(x)
        # Apply mask
        out = out * mask.unsqueeze(-1)
        # Get CRF output
        out = self.crf.decode(out, mask.bool())
        return out

    def log_likelihood(self, x: Tensor, tags: Tensor, mask: Tensor) -> Tensor:
        # Get LSTM output
        out = self.get_lstm2linear(x)
        if isnan(out).any() or isinf(out).any():
            return tensor(0.0, requires_grad=True)
        # Apply mask
        out = out * mask.unsqueeze(-1)
        # Get CRF output
        return -self.crf(
            out, tags,
            mask.bool(),
            reduction='mean'
        )
